{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"docs/images/DSPy8.png\" alt=\"DSPy7 Image\" height=\"150\"/>\n",
    "\n",
    "# DSPy: Tutorial @ SkyCamp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook contains the **DSPy tutorial** for **SkyCamp 2023**.\n",
    "\n",
    "Let's begin by setting things up. The snippet below will also install **DSPy** if it's not there already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "import os\n",
    "\n",
    "try: # When on google Colab, let's clone the notebook so we download the cache.\n",
    "    import google.colab\n",
    "    repo_path = 'dspy'\n",
    "    !git -C $repo_path pull origin || git clone https://github.com/stanfordnlp/dspy $repo_path\n",
    "except:\n",
    "    repo_path = '.'\n",
    "\n",
    "if repo_path not in sys.path:\n",
    "    sys.path.append(repo_path)\n",
    "\n",
    "# Set up the cache for this notebook\n",
    "os.environ[\"DSP_NOTEBOOK_CACHEDIR\"] = os.path.join(repo_path, 'cache')\n",
    "\n",
    "import pkg_resources # Install the package if it's not installed\n",
    "if not \"dspy-ai\" in {pkg.key for pkg in pkg_resources.working_set}:\n",
    "    !pip install -U pip\n",
    "    !pip install dspy-ai==2.1\n",
    "    # !pip install -e $repo_path\n",
    "\n",
    "!pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dspy\n",
    "from dspy.evaluate import Evaluate\n",
    "from dspy.teleprompt import BootstrapFewShot, BootstrapFewShotWithRandomSearch, BootstrapFinetune"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1) Configure the default LM and retriever\n",
    "\n",
    "We'll start by setting up the language model (LM) and retrieval model (RM). **DSPy** supports multiple API and local models.\n",
    "\n",
    "In this notebook, we will use `Llama2-13b-chat` using the HuggingFace TGI serving software infrastructure. In principle you can run this on your own local GPUs, but for this tutorial all examples are pre-cached so you don't need to worry about cost.\n",
    "\n",
    "We will use the retriever `ColBERTv2`. To make things easy, we've set up a ColBERTv2 server hosting a Wikipedia 2017 \"abstracts\" search index (i.e., containing first paragraph of each article from this [2017 dump](https://hotpotqa.github.io/wiki-readme.html)), so you don't need to worry about setting one up! It's free.\n",
    "\n",
    "**Note:** _If you run this notebook as instructed, you don't need an API key. All examples are already cached internally so you can inspect them!_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama = dspy.HFClientTGI(model=\"meta-llama/Llama-2-13b-chat-hf\", port=[7140, 7141, 7142, 7143], max_tokens=150)\n",
    "colbertv2 = dspy.ColBERTv2(url='http://20.102.90.50:2017/wiki17_abstracts')\n",
    "\n",
    "# # NOTE: After you finish this notebook, you can use GPT-3.5 like this if you like.\n",
    "# turbo = dspy.OpenAI(model='gpt-3.5-turbo-instruct')\n",
    "# # In that case, make sure to configure lm=turbo below if you choose to do that.\n",
    "\n",
    "dspy.settings.configure(rm=colbertv2, lm=llama)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2) Create a few question–answer pairs for our task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = [('Who was the director of the 2009 movie featuring Peter Outerbridge as William Easton?', 'Kevin Greutert'),\n",
    "         ('The heir to the Du Pont family fortune sponsored what wrestling team?', 'Foxcatcher'),\n",
    "         ('In what year was the star of To Hell and Back born?', '1925'),\n",
    "         ('Which award did the first book of Gary Zukav receive?', 'U.S. National Book Award'),\n",
    "         ('What documentary about the Gilgo Beach Killer debuted on A&E?', 'The Killing Season'),\n",
    "         ('Which author is English: John Braine or Studs Terkel?', 'John Braine'),\n",
    "         ('Who produced the album that included a re-recording of \"Lithium\"?', 'Butch Vig')]\n",
    "\n",
    "train = [dspy.Example(question=question, answer=answer).with_inputs('question') for question, answer in train]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev = [('Who has a broader scope of profession: E. L. Doctorow or Julia Peterkin?', 'E. L. Doctorow'),\n",
    "       ('Right Back At It Again contains lyrics co-written by the singer born in what city?', 'Gainesville, Florida'),\n",
    "       ('What year was the party of the winner of the 1971 San Francisco mayoral election founded?', '1828'),\n",
    "       ('Anthony Dirrell is the brother of which super middleweight title holder?', 'Andre Dirrell'),\n",
    "       ('The sports nutrition business established by Oliver Cookson is based in which county in the UK?', 'Cheshire'),\n",
    "       ('Find the birth date of the actor who played roles in First Wives Club and Searching for the Elephant.', 'February 13, 1980'),\n",
    "       ('Kyle Moran was born in the town on what river?', 'Castletown River'),\n",
    "       (\"The actress who played the niece in the Priest film was born in what city, country?\", 'Surrey, England'),\n",
    "       ('Name the movie in which the daughter of Noel Harrison plays Violet Trefusis.', 'Portrait of a Marriage'),\n",
    "       ('What year was the father of the Princes in the Tower born?', '1442'),\n",
    "       ('What river is near the Crichton Collegiate Church?', 'the River Tyne'),\n",
    "       ('Who purchased the team Michael Schumacher raced for in the 1995 Monaco Grand Prix in 2000?', 'Renault'),\n",
    "       ('André Zucca was a French photographer who worked with a German propaganda magazine published by what Nazi organization?', 'the Wehrmacht')]\n",
    "\n",
    "dev = [dspy.Example(question=question, answer=answer).with_inputs('question') for question, answer in dev]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3) Key Concepts: Signatures & Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a dspy.Predict module with the signature `question -> answer` (i.e., takes a question and outputs an answer).\n",
    "predict = dspy.Predict('question -> answer')\n",
    "\n",
    "# Use the module!\n",
    "predict(question=\"What is the capital of Germany?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the example above, we used the `dspy.Predict` module **zero-shot**, i.e. without compiling it on any examples.\n",
    "\n",
    "Let's now build a slightly more advanced program. Our program will use the `dspy.ChainOfThought` module, which asks the LM to think step by step.\n",
    "\n",
    "We will call this program `CoT`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CoT(dspy.Module):  # let's define a new module\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        # here we declare the chain of thought sub-module, so we can later compile it (e.g., teach it a prompt)\n",
    "        self.generate_answer = dspy.ChainOfThought('question -> answer')\n",
    "    \n",
    "    def forward(self, question):\n",
    "        return self.generate_answer(question=question)  # here we use the module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's compile this using our seven `train` examples. We will use the very simple `BootstrapFewShot` in DSPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_EM = dspy.evaluate.answer_exact_match\n",
    "\n",
    "teleprompter = BootstrapFewShot(metric=metric_EM, max_bootstrapped_demos=2)\n",
    "cot_compiled = teleprompter.compile(CoT(), trainset=train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's ask a question to this new program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cot_compiled(\"What is the capital of Germany?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might be curious what's happening under the hood. Let's inspect the last call to our Llama LM to see the prompt and the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the prompt ends with the question we asked (\"What is the capital of Germany?\"), but before that it includes few-shot examples.\n",
    "\n",
    "The final example in the prompt contains a rationale (step-by-step reasoning) self-generated from the LM for use as a demonstration, for the training question \"Which author is English: John Braine or Studs Terkel?\".\n",
    "\n",
    "Now, let's evaluate on our development set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_THREADS = 32\n",
    "evaluate_hotpot = Evaluate(devset=dev, metric=metric_EM, num_threads=NUM_THREADS, display_progress=True, display_table=15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's evaluate the compiled `CoT` program with Llama. Feel free to replace `cot_compiled` below with `CoT()` (notice the paranthesis) to test the zero-shot version of CoT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_hotpot(cot_compiled)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4) Bonus 1: RAG with query generation\n",
    "\n",
    "As a bonus, let's define a more sophisticated program called `RAG`. This program will:\n",
    "\n",
    "- Use the LM to generate a search query based on the input question\n",
    "- Retrieve three passages using our retriever\n",
    "- Use the LM to generate a final answer using these passages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RAG(dspy.Module):\n",
    "    def __init__(self, num_passages=3):\n",
    "        super().__init__()\n",
    "\n",
    "        # declare three modules: the retriever, a query generator, and an answer generator\n",
    "        self.retrieve = dspy.Retrieve(k=num_passages)\n",
    "        self.generate_query = dspy.ChainOfThought(\"question -> search_query\")\n",
    "        self.generate_answer = dspy.ChainOfThought(\"context, question -> answer\")\n",
    "    \n",
    "    def forward(self, question):\n",
    "        # generate a search query from the question, and use it to retrieve passages\n",
    "        search_query = self.generate_query(question=question).search_query\n",
    "        passages = self.retrieve(search_query).passages\n",
    "\n",
    "        # generate an answer from the passages and the question\n",
    "        return self.generate_answer(context=passages, question=question)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Out of curiosity, we can evaluate the **uncompiled** (or **zero-shot**) version of this program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_hotpot(RAG(), display_table=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now compile this RAG program. We'll use a slightly more advanced teleprompter (automatic prompt optimizer) this time, which relies on random search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "teleprompter2 = BootstrapFewShotWithRandomSearch(metric=metric_EM, max_bootstrapped_demos=2, num_candidate_programs=8, num_threads=NUM_THREADS)\n",
    "rag_compiled = teleprompter2.compile(RAG(), trainset=train, valset=dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now evaluate this compiled version of RAG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_hotpot(rag_compiled)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect one of the LM calls for this. Focus in particular on the structure of the last few input/output examples in the prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rag_compiled(\"What year was the party of the winner of the 1971 San Francisco mayoral election founded?\")\n",
    "llama.inspect_history(n=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4) Bonus 2: Multi-Hop Retrieval and Reasoning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now build a simple multi-hop program, which will interleave multiple calls to the LM and the retriever.\n",
    "\n",
    "Please follow the **TODO** instructions below to implement this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsp.utils.utils import deduplicate\n",
    "\n",
    "class MultiHop(dspy.Module):\n",
    "    def __init__(self, num_passages=3):\n",
    "        super().__init__()\n",
    "\n",
    "        self.retrieve = dspy.Retrieve(k=num_passages)\n",
    "        self.generate_query = dspy.ChainOfThought(\"question -> search_query\")\n",
    "\n",
    "        # TODO: Define a dspy.ChainOfThought module with the signature 'context, question -> search_query'.\n",
    "        self.generate_query_from_context = None\n",
    "\n",
    "        self.generate_answer = dspy.ChainOfThought(\"context, question -> answer\")\n",
    "    \n",
    "    def forward(self, question):\n",
    "        passages = []\n",
    "        \n",
    "        search_query = self.generate_query(question=question).search_query\n",
    "        passages += self.retrieve(search_query).passages\n",
    "\n",
    "        # TODO: Replace `None` with a call to self.generate_query_from_context to generate a search query.\n",
    "        # Note: In DSPy, always pass keyword arguments (e.g., context=..., question=...) to the modules to avoid ambiguity.\n",
    "        # Note 2: Don't forget to access the field .search_query to extract that from the output of the module.\n",
    "        # Note 3: Check the following notebook for a completed example: https://github.com/stanfordnlp/dspy/blob/main/skycamp2023_completed.ipynb.\n",
    "        search_query2 = None\n",
    "\n",
    "        # TODO: Replace `None` with a call to self.retrieve to retrieve passages. Append them to the list `passages`.\n",
    "        passages += None\n",
    "\n",
    "        return self.generate_answer(context=deduplicate(passages), question=question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "multihop_compiled = teleprompter2.compile(MultiHop(), trainset=train, valset=dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate_hotpot(multihop_compiled, devset=dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now inspect the prompt for the second-hop search query for one of the questions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "multihop_compiled(question=\"Who purchased the team Michael Schumacher raced for in the 1995 Monaco Grand Prix in 2000?\")\n",
    "llama.inspect_history(n=1, skip=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py39_aug2023_dspy",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
